import gradio as gr
import pickle
import numpy as np
import cv2  # 导入OpenCV库

# 加载保存的KNN模型
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)

def convert_to_grayscale(img_array):
    if img_array.ndim == 3 and img_array.shape[2] == 3:
        grayscale = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
    elif img_array.ndim == 3 and img_array.shape[2] == 4:
        # 只取RGB通道，忽略透明度通道
        img_array = img_array[..., :3]
        grayscale = cv2.cvtColor(img_array, cv2.COLOR_BGR2GRAY)
    else:
        grayscale = img_array
    return grayscale

def preprocess_image(img):
    if img.ndim == 3 and img.shape[2] == 4:
        # 只取RGB通道，忽略透明度通道
        img_array = img[..., :3] * (img[..., 3:4] / 255.0)
    else:
        img_array = img

    if img_array.ndim != 2 and img_array.ndim != 3:
        raise ValueError("图像必须是2D（灰度图）或3D（彩色图).")

    img_array = img_array.astype(np.uint8)
    
    img_resized = cv2.resize(img_array, (8, 8), interpolation=cv2.INTER_CUBIC)
    img_array = convert_to_grayscale(img_resized)

    img_array = img_array.reshape(-1) / 255.0
    img_array = img_array.astype(np.float32)

    return img_array

def predict_digit(sketch):
    if sketch.max() == 0:
        return "未检测到输入，请绘制一个数字。"
    
    img_data = sketch
    img = preprocess_image(img_data)
    img = np.expand_dims(img, axis=0)

    prediction = knn_model.predict(img)[0]
    return str(prediction)

# 创建Gradio界面
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(label="在这里绘制你的数字", type="numpy"),  # 更改这里
    outputs=gr.Label(label="预测"),
    title="手写数字识别",
    description="在画布上绘制你的数字并获取预测结果。"
)

# 启动Gradio界面
iface.launch()